/** @file pegasos.h
 *
 *  @author Dongryeol Lee (drselee@gmail.com)
 */

#ifndef MLPACK_SVM_PEGASOS_H
#define MLPACK_SVM_PEGASOS_H

#include "core/util/random.h"
#include "mlpack/svm/svm_arguments.h"
#include "mlpack/svm/svm_model.h"
#include "mlpack/svm/svm_result.h"

namespace mlpack {
namespace svm {

class Pegasos {
  public:

    static bool Train(
      const core::abstract::AbstractArguments *arguments_in,
      core::abstract::AbstractResult *result_out,
      core::abstract::AbstractModel *model_out)  {

      // Guaranteed to be an SVMArguments, SVMResult, SVMModel here.
      const mlpack::svm::SVMArguments *svm_arguments =
        dynamic_cast< const mlpack::svm::SVMArguments *>(arguments_in);
      if(! svm_arguments) {
        std::cerr <<
                  "You need to call this function with an " <<
                  "SVMArgument object.\n";
        return false;
      }
      mlpack::svm::SVMResult *svm_result =
        dynamic_cast< mlpack::svm::SVMResult *>(result_out);
      if(! svm_result) {
        std::cerr <<
                  "You need to call this function with a " <<
                  "SVMResult object.\n";
      }
      mlpack::svm::SVMModel *svm_model =
        dynamic_cast< mlpack::svm::SVMModel *>(model_out);
      if(! svm_model) {
        std::cerr <<
                  "You need to call this function with a " <<
                  "SVMModel object.\n";
      }

      // Initialize the weight vector.
      svm_result->weight_vector_.zeros(
        svm_arguments->reference_table_->num_point_attributes());

      // The random number generator.
      core::util::Random random;

      // The contribution from the mistakes.
      arma::vec mistake_contribution;

      for(int i = 1; i <= svm_arguments->max_num_iterations_ ; i++) {

        // Determine the learning rate.
        double eta = 1.0 / (svm_arguments->lambda_ * i);

        // Randomly choose a point set and count the percentage of
        // mistakes.
        mistake_contribution.zeros(
          svm_arguments->reference_table_->num_point_attributes());

        for(int j = 0;
            j < svm_arguments->num_random_train_samples_per_iter_; j++) {

          // Select a random point and compute loss.
          int random_index =
            random.Integer(svm_arguments->reference_table_->num_points());
          const core::data::alias_vec *random_instance_target = NULL;
          const core::data::alias_vec &random_instance =
            svm_arguments->reference_table_->get_instance(
              random_index, &random_instance_target);
          double prediction =
            arma::dot(svm_result->weight_vector_, random_instance);
          int label = static_cast<int>((*random_instance_target)[0]);
          double loss =
            std::max(1.0 - label * prediction, 0.0);

          if(loss > 0.0) {
            mistake_contribution *= j;
            arma::scaled_add(label, random_instance, &mistake_contribution);
            mistake_contribution /= static_cast<double>(j + 1);
          }
        }

        // Scale the previous weight and add mistake contributions.
        svm_result->weight_vector_ *= (1 - eta * svm_arguments->lambda_);
        svm_result->weight_vector_ += eta * mistake_contribution;

        // Project.
        double weight_vector_squared_norm =
          arma::dot(svm_result->weight_vector_, svm_result->weight_vector_);
        if(weight_vector_squared_norm > 1.0 / svm_arguments->lambda_) {
          svm_result->weight_vector_ *=
            sqrt(
              1.0 / (svm_arguments->lambda_ * weight_vector_squared_norm)) ;
        }
      } // end of looping over each iteration.

      std::cerr << "Trained Pegasos SVM weight vector...\n";
      svm_result->weight_vector_.print();

      return true;
    }
};
}
}

#endif
